import os.path
from typing import Optional, Callable

import pandas as pd
import torch
from PIL import Image
from torch.utils.data import Dataset


class Hotdog(Dataset):
    def __init__(self,
                 root_path: str,
                 train: bool,
                 transform: Optional[Callable] = None):
        self.root_path = root_path
        self.train = train
        self.transform = transform
        self.df = pd.read_csv(os.path.join(root_path, 'train.csv' if train else 'test.csv'))

    def __getitem__(self, index):
        q = self.df.loc[index]
        path = os.path.join(self.root_path, 'train' if self.train else 'test', q['path'])
        label = q['label']
        img = Image.open(path)
        if self.transform is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.df)
